import tkinter as tk
from tkinter import messagebox
from openai import OpenAI


class ChatApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Chat with OpenAI")

        self.message_frame = tk.Frame(root)
        self.message_frame.pack(padx=10, pady=10)

        self.messages_text = tk.Text(self.message_frame, wrap="word", state="disabled")
        self.messages_text.pack(expand=True, fill="both")

        self.input_frame = tk.Frame(root)
        self.input_frame.pack(padx=10, pady=(0, 10))

        self.input_entry = tk.Entry(self.input_frame)
        self.input_entry.pack(side="left", expand=True, fill="x")
        self.input_entry.bind("<Return>", lambda event: self.send_message())

        self.send_button = tk.Button(self.input_frame, text="发送", command=self.send_message)
        self.send_button.pack(side="right")

        self.client = OpenAI(
            base_url="https://integrate.api.nvidia.com/v1",
            api_key="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"
        )

        self.messages_all = []

    def send_message(self):
        user_input = self.input_entry.get().strip()
        if user_input:
            self.messages_all.append({
                "content": user_input,
                "role": "user"
            })
            completion = self.client.chat.completions.create(
                model="google/codegemma-7b",
                messages=self.messages_all,
                temperature=0.5,
                top_p=1,
                max_tokens=1024,
                stream=True
            )
            response_msg = ""
            self.messages_text.config(state="normal")
            self.insert_message_to_display(f"You: {user_input}", "user")
            for chunk in completion:
                if chunk.choices[0].delta.content is not None:
                    temp_msg = chunk.choices[0].delta.content
                    response_msg += temp_msg
            self.insert_message_to_display(f"Assistant: {response_msg}", "assistant")
            self.messages_text.config(state="disabled")
            self.messages_text.see("end")
            self.input_entry.delete(0, "end")

    def insert_message_to_display(self, message, role):
        # 设置对齐方式
        align = "right" if role == "user" else "left"
        # 在Text组件中插入消息
        self.messages_text.insert("end", f"{message}\n", role)
        # 设置插入的消息的对齐方式
        self.messages_text.tag_configure(role, justify=align)


root = tk.Tk()
app = ChatApp(root)
root.mainloop()
